#!/usr/bin/env python

import asyncio
import datetime
import random
import websockets
import json
import time

from usb_camera import Usb_Camera
from file_io import File_IO

class ROS_Server:
    def __init__(self):
        self.func_map = {}
        self.websocket = None

    def add_handler(self, handler):
        hdl = handler()
        hdl.register(self.func_map)

    async def call_service(self, service, args, id):
        if service in self.func_map:
            # print(service, args)
            await self.func_map[service](self.websocket, args, id)
        else:
            print('unhandled service: ', service)



connected_clients = set()

async def broadcast_message(message):
    # 移除已经关闭的连接
    disconnected_clients = {client for client in connected_clients if client.closed}
    for client in disconnected_clients:
        connected_clients.remove(client)

    # 向所有活跃的客户端发送消息
    for client in connected_clients:
        await client.send(message)

async def ws_handler(websocket, path):
    # async for message in websocket:
    #     message = "got message: {}".format(message)
    print('client connectted')
    # 客户端加入集合
    connected_clients.add(websocket)
    ros_server = ROS_Server()
    ros_server.add_handler(Usb_Camera)
    ros_server.add_handler(File_IO)
    ros_server.websocket = websocket
    while True:
        try:
            recv_text = await websocket.recv()
            # print(recv_text)
            msg = json.loads(recv_text)
            print(msg['op'], msg['id'])
            if msg['op'] == 'call_service' and 'id' in msg:
                params = msg['id'].split(':')
                service = params[1]
                id = params[2]
                args = msg['args']
                await ros_server.call_service(service, args, id)
            elif msg['op'] == 'publish':
                await broadcast_message(json.dumps({"op":msg['op'], "topic": msg['topic'], "msg":msg['msg']}))
                # await websocket.send()
            else:
                print(msg)
        except Exception as e:
            break
            print(e)

start_server = websockets.serve(ws_handler, '0.0.0.0', 9090)

asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()